import asyncio
import ssl
import json
import websockets
from urllib.parse import urlparse, parse_qs


class ChatRoom:
    def __init__(self, room_id):
        self.room_id = room_id
        self.users = set()
        self.history = []

    async def notify_users(self):
        """通知当前聊天室的用户数量"""
        if self.users:
            message = json.dumps({"type": "user_count", "count": len(self.users)})
            tasks = [asyncio.create_task(user.send(message)) for user in self.users]
            await asyncio.wait(tasks)

    async def send_message_history(self, websocket):
        """将历史消息发送给新连接的用户"""
        for message in self.history:
            await websocket.send(message)

    async def broadcast_message(self, message, websocket):
        """广播消息给当前聊天室的用户"""
        tasks = [asyncio.create_task(user.send(message)) for user in self.users
                 # if websocket is None or user != websocket
                 ]
        if tasks:  # 如果任务列表不为空
            await asyncio.wait(tasks)

    def cache_history(self, message):
        """添加到历史消息缓存并管理消息"""
        self.history.append(message)

    def clear_history(self):
        """将历史消息保存到文件并清空"""
        with open(f'chat_history/{self.room_id}', 'a', encoding='utf-8') as f:
            for msg in self.history:
                f.write(msg + '\n')
        self.history.clear()

    def manage_history(self):
        """管理历史消息，超出部分保存到文件中"""
        if len(self.history) > 2000:
            old_messages = self.history[:-1500]
            self.history = self.history[-1500:]
            with open(f'chat_history/{self.room_id}', 'a', encoding='utf-8') as f:
                for msg in old_messages:
                    f.write(msg + '\n')


class ChatServer:
    def __init__(self):
        self.chat_rooms = {}

    def get_room(self, room_id):
        """获取或创建一个聊天室"""
        if room_id not in self.chat_rooms:
            self.chat_rooms[room_id] = ChatRoom(room_id)
        return self.chat_rooms[room_id]

    async def chat_handler(self, websocket, path):
        # 解析URL，获取聊天室ID
        query_params = parse_qs(urlparse(path).query)
        room_id = query_params.get('c', ['default'])[0]

        room = self.get_room(room_id)
        room.users.add(websocket)
        await room.send_message_history(websocket)
        await room.notify_users()

        try:
            async for message in websocket:
                if "clear" in message:
                    room.clear_history()
                    broadcast_message = json.dumps({"type": "message", "content": "CLEARED."})
                    room.cache_history(broadcast_message)
                    await room.broadcast_message(broadcast_message, None)
                else:
                    room.manage_history()
                    broadcast_message = json.dumps({"type": "message", "content": message})
                    room.cache_history(broadcast_message)
                    await room.broadcast_message(broadcast_message, websocket)
        finally:
            room.users.remove(websocket)
            await room.notify_users()


async def start_server():
    chat_server = ChatServer()

    try:
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ssl_context.load_cert_chain(certfile='/etc/letsencrypt/live/www.hei.asia/fullchain.pem',
                                    keyfile='/etc/letsencrypt/live/www.hei.asia/privkey.pem')

        server = websockets.serve(chat_server.chat_handler, '0.0.0.0', 8765, ssl=ssl_context)
        print("Running with SSL")
    except Exception as e:
        print(f"Error: {e}")
        server = websockets.serve(chat_server.chat_handler, '0.0.0.0', 8765)
        print("Running without SSL")

    await server


if __name__ == "__main__":
    asyncio.get_event_loop().run_until_complete(start_server())
    asyncio.get_event_loop().run_forever()
